# -*- coding: utf-8 -*-
"""
分布拟合模块

提供正态分布和指数分布的拟合检验功能，包括：
- normal_fit_plot: 正态分布拟合与Shapiro-Wilk检验
- exponential_fit_plot: 指数分布拟合与KS检验
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats

try:
    # 用于包内调用
    from .draw_method import draw_distribution_fit
except ImportError:
    # 用于直接运行
    from draw_method import draw_distribution_fit

__all__ = ['normal_fit_plot', 'exponential_fit_plot']

def normal_fit_plot(data, save_path=None, alpha_level = 0.05, bins=30, alpha=0.6, show=True, title='正态分布拟合与检验', label_data='数据分布', label_fit='拟合分布'):
	"""
	对数据进行正态性检验，并拟合正态分布，绘制直方图和拟合曲线。
	参数：
		data: 一维array-like，待检验和拟合的数据
		save_path: str，可选，保存图片的路径
		bins: int，直方图分箱数
		alpha: float，直方图透明度
		alpha_level: float，大于该值认为假设可信
		show: bool，是否显示图像
		title: str，图标题
		label_data, label_fit: str，图例标签
	返回:
        如果不满足假设返回 None, 否则返回:
		检验结果字典，包括shapiro统计量、p值、拟合参数等
	"""
	data = np.asarray(data)
	# 正态性检验
	shapiro_stat, shapiro_p = stats.shapiro(data)
	normal = shapiro_p > alpha_level

	if not normal:
		print(f"拒绝 {alpha_level} 下的正态分布假设! (p={shapiro_p:.4f})")
		return None

	# 拟合参数
	mu, sigma = stats.norm.fit(data)

	# 可视化 - 使用draw_method中的函数
	if save_path:
		data_series = pd.Series(data)
		draw_distribution_fit(
			data=data_series,
			fit_func=stats.norm.pdf,
			fit_params=(mu, sigma),
			title=title,
			save_path=save_path,
			bins=bins,
			alpha=alpha,
			label_data=label_data,
			label_fit=label_fit
		)
	elif show:
		# 如果没有保存路径但需要显示，使用临时路径
		data_series = pd.Series(data)
		draw_distribution_fit(
			data=data_series,
			fit_func=stats.norm.pdf,
			fit_params=(mu, sigma),
			title=title,
			save_path='temp_normal_fit.png',
			bins=bins,
			alpha=alpha,
			label_data=label_data,
			label_fit=label_fit
		)
	
	return {
		'shapiro_stat': shapiro_stat,
		'shapiro_p': shapiro_p,
		'is_normal': normal,
		'mu': mu,
		'sigma': sigma
	}

def exponential_fit_plot(data, save_path=None, alpha_level=0.05, bins=30, alpha=0.6, show=True, title='指数分布拟合与检验', label_data='数据分布', label_fit='拟合分布'):
	"""
	对数据进行指数分布检验，并拟合指数分布，绘制直方图和拟合曲线。
	参数：
		data: 一维array-like，待检验和拟合的数据（必须为非负值）
		save_path: str，可选，保存图片的路径
		bins: int，直方图分箱数
		alpha: float，直方图透明度
		alpha_level: float，大于该值认为假设可信
		show: bool，是否显示图像
		title: str，图标题
		label_data, label_fit: str，图例标签
	返回:
        如果不满足假设返回 None, 否则返回:
		检验结果字典，包括KS统计量、p值、拟合参数等
	"""
	data = np.asarray(data)
	
	# 检查数据是否为非负值
	if np.any(data < 0):
		print("警告：指数分布要求数据为非负值，已自动过滤负值")
		data = data[data >= 0]
	
	if len(data) == 0:
		print("错误：过滤后无有效数据")
		return None
	
	# 拟合指数分布参数
	# scipy.stats.expon.fit() 返回 (loc, scale)，其中 loc 通常为0，scale为1/lambda
	loc, scale = stats.expon.fit(data)
	
	# KS检验指数分布拟合度
	ks_stat, ks_p = stats.kstest(data, lambda x: stats.expon.cdf(x, loc=loc, scale=scale))
	is_exponential = ks_p > alpha_level
	
	if not is_exponential:
		print(f"拒绝 {alpha_level} 下的指数分布假设! (p={ks_p:.4f})")
		return None
	
	# 可视化 - 使用draw_method中的函数
	if save_path:
		data_series = pd.Series(data)
		draw_distribution_fit(
			data=data_series,
			fit_func=stats.expon.pdf,
			fit_params=(loc, scale),
			title=title,
			save_path=save_path,
			bins=bins,
			alpha=alpha,
			label_data=label_data,
			label_fit=label_fit
		)
	elif show:
		# 如果没有保存路径但需要显示，使用临时路径
		data_series = pd.Series(data)
		draw_distribution_fit(
			data=data_series,
			fit_func=stats.expon.pdf,
			fit_params=(loc, scale),
			title=title,
			save_path='temp_exponential_fit.png',
			bins=bins,
			alpha=alpha,
			label_data=label_data,
			label_fit=label_fit
		)
	
	return {
		'ks_stat': ks_stat,
		'ks_p': ks_p,
		'is_exponential': is_exponential,
		'loc': loc,
		'scale': scale,
		'lambda': 1/scale  # 指数分布的率参数
	}

# 测试代码
if __name__ == "__main__":
	np.random.seed(42)
	
	# 测试正态分布拟合
	print("="*50)
	print("正态分布拟合测试")
	print("="*50)
	true_mu, true_sigma = 5.0, 2.0
	normal_data = stats.norm.rvs(loc=true_mu, scale=true_sigma, size=1000)
	result = normal_fit_plot(normal_data, save_path='test_normal_fit.png', show=True)
	if result is not None:
		print(f"Shapiro-Wilk检验统计量: {result['shapiro_stat']:.4f}, p值: {result['shapiro_p']:.4f}")
		print(f"拟合均值: {result['mu']:.4f}, 拟合标准差: {result['sigma']:.4f}")
	else:
		print("数据未通过正态性检验，未能输出统计量和p值。")
	
	# 测试指数分布拟合
	print("\n" + "="*50)
	print("指数分布拟合测试")
	print("="*50)
	true_lambda = 0.5  # 率参数
	exponential_data = stats.expon.rvs(scale=1/true_lambda, size=1000)
	result_exp = exponential_fit_plot(exponential_data, save_path='test_exponential_fit.png', show=True)
	if result_exp is not None:
		print(f"KS检验统计量: {result_exp['ks_stat']:.4f}, p值: {result_exp['ks_p']:.4f}")
		print(f"拟合位置参数: {result_exp['loc']:.4f}, 拟合尺度参数: {result_exp['scale']:.4f}")
		print(f"拟合率参数λ: {result_exp['lambda']:.4f}")
	else:
		print("数据未通过指数分布检验，未能输出统计量和p值。")
